#!/usr/bin/env python3
"""
ESP32摄像头ROS2节点
适配echokitbot项目
支持UDP自动发现和HTTP流式传输

功能特性：
- UDP自动发现ESP32-CAM设备
- HTTP MJPEG流解析
- 动态参数配置和同步
- 自动重连机制
- 帧率统计和性能监控
- 诊断信息发布
- 服务接口支持
"""

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage
from std_srvs.srv import Trigger
from diagnostic_msgs.msg import DiagnosticArray, DiagnosticStatus, KeyValue
from rcl_interfaces.msg import ParameterDescriptor, IntegerRange
from rcl_interfaces.msg import SetParametersResult
import urllib.error
import urllib.request
import time
import threading
import socket
from enum import Enum
from queue import Queue, Full


class CameraState(Enum):
    """摄像头状态枚举"""
    DISCONNECTED = 0
    DISCOVERING = 1
    CONNECTING = 2
    CONNECTED = 3
    ERROR = 4


class ESP32CameraNode(Node):
    """ESP32摄像头ROS2节点"""
    
    def __init__(self):
        super().__init__('esp32_camera_node')
        
        # 发布器 - 发布压缩图像
        self.publisher_ = self.create_publisher(CompressedImage, 'camera/image_raw/compressed', 50)
        self.diagnostics_pub = self.create_publisher(
            DiagnosticArray, '/diagnostics', 10
        )
        
        # 状态变量
        self.camera_state = CameraState.DISCOVERING
        self.last_recv_img_time = 0
        self.camera_ip = None
        self.stream = None
        self.stream_bytes = b''
        self.count = 0
        self.last_time = time.time()
        self.rate = 0
        self.total_frames = 0
        self.dropped_frames = 0
        self.reconnect_count = 0
        self.last_error = None
        self.skipped_frames = 0
        
        # 图像处理队列 - 增大队列
        self.image_queue = Queue(maxsize=20)
        self.processing_thread = None
        
        # 线程锁
        self.state_lock = threading.Lock()
        self.stream_thread = None
        
        # 声明参数
        self._declare_camera_parameters()
        self._declare_node_parameters()
        
        # 添加参数回调
        self.add_on_set_parameters_callback(self.parameter_callback)
        
        # 创建服务
        self.reconnect_service = self.create_service(
            Trigger, 'reconnect_camera', self.reconnect_callback
        )
        self.get_status_service = self.create_service(
            Trigger, 'get_camera_status', self.get_status_callback
        )
        
        # 创建定时器发布诊断信息
        self.diagnostics_timer = self.create_timer(1.0, self.publish_diagnostics)
        
        # 启动图像处理线程
        self.processing_thread = threading.Thread(target=self.process_images, daemon=True)
        self.processing_thread.start()
        
        # 启动UDP监听线程
        self.thread_wait_connect = threading.Thread(target=self.wait_connect, daemon=True)
        self.thread_wait_connect.start()
        
        self.get_logger().info('ESP32摄像头节点已启动，等待摄像头连接...')
        self.get_logger().info('优化配置: 320x240 | JPEG质量18 | 压缩图像传输 | 队列20帧')
        self.get_logger().info('服务接口:')
        self.get_logger().info('  - ~/reconnect_camera: 手动触发重连')
        self.get_logger().info('  - ~/get_camera_status: 获取摄像头状态')
    
    def _declare_camera_parameters(self):
        """声明摄像头参数"""
        # 图像质量 (4-64, 数值越小质量越高)
        self.declare_param_with_range('quality', 18, 4, 64, 
            'JPEG图像质量 (4-64, 数值越小质量越高)')
        
        # 图像大小 (0-13)
        # 0: 96x96, 5: 320x240, 8: 640x480, 10: 800x600, 13: 1600x1200
        self.declare_param_with_range('framesize', 5, 0, 13,
            '图像尺寸 (0:96x96, 5:320x240, 8:640x480, 10:800x600, 13:1600x1200)')
        
        # LED亮度 (0-255)
        self.declare_param_with_range('led_intensity', 0, 0, 255,
            'LED亮度 (0-255)')
        
        # 水平镜像 (0: 关闭, 1: 开启)
        self.declare_param_with_range('hmirror', 0, 0, 1,
            '水平镜像 (0:关闭, 1:开启)')
        
        # 垂直翻转 (0: 关闭, 1: 开启)
        self.declare_param_with_range('vflip', 0, 0, 1,
            '垂直翻转 (0:关闭, 1:开启)')
        
        # 人脸检测 (0: 关闭, 1: 开启)
        self.declare_param_with_range('face_detect', 0, 0, 1,
            '人脸检测 (0:关闭, 1:开启)')
        
        # 特效 (0-6)
        # 0: 无, 1: 负片, 2: 灰度, 3: 红色, 4: 绿色, 5: 蓝色, 6: 复古
        self.declare_param_with_range('special_effect', 0, 0, 6,
            '特效 (0:无, 1:负片, 2:灰度, 3:红色, 4:绿色, 5:蓝色, 6:复古)')
    
    def _declare_node_parameters(self):
        """声明节点参数"""
        self.declare_parameter('udp_port', 8887,
            ParameterDescriptor(description='UDP发现端口'))
        self.declare_parameter('stream_port', 81,
            ParameterDescriptor(description='HTTP流端口'))
        self.declare_parameter('reconnect_timeout', 30.0,
            ParameterDescriptor(description='重连超时时间(秒)'))
        self.declare_parameter('stream_timeout', 10.0,
            ParameterDescriptor(description='流连接超时时间(秒)'))
        self.declare_parameter('enable_diagnostics', True,
            ParameterDescriptor(description='启用诊断信息发布'))
        self.declare_parameter('enable_frame_skip', True,
            ParameterDescriptor(description='启用帧跳过（队列满时丢弃旧帧）'))
        self.declare_parameter('max_buffer_size', 5242880,
            ParameterDescriptor(description='最大缓冲区大小(字节)，防止内存溢出'))
        self.declare_parameter('read_chunk_size', 16384,
            ParameterDescriptor(description='每次读取的数据块大小(字节)'))
    
    def declare_param_with_range(self, name: str, value: int, start: int, end: int, 
                                  description: str = ''):
        """声明带范围的整数参数"""
        pd = ParameterDescriptor(description=description)
        pd_range = IntegerRange()
        pd_range.from_value = start
        pd_range.to_value = end
        pd_range.step = 1
        pd.integer_range.append(pd_range)
        self.declare_parameter(name, value, pd)
    
    def parameter_callback(self, parameters):
        """参数变更回调"""
        for parameter in parameters:
            self.get_logger().info(f'参数 {parameter.name} 设置为：{parameter.value}')
            self.sync_param(parameter.name, parameter.value)
        return SetParametersResult(successful=True)
    
    def sync_param(self, name: str, value: int) -> bool:
        """同步参数到摄像头
        
        Returns:
            bool: 同步是否成功
        """
        if not self.camera_ip:
            self.get_logger().warn(f'摄像头未连接，无法同步参数 {name}')
            return False
            
        url = f'http://{self.camera_ip}/control?var={name}&val={value}'
        try:
            with urllib.request.urlopen(url, timeout=2) as response:
                result = response.read().decode('utf-8')
                self.get_logger().debug(f'参数同步成功: {name}={value}, 响应: {result}')
                return True
        except urllib.error.URLError as e:
            self.get_logger().error(f"参数同步失败 {name}={value}: {e}")
            self.last_error = f"参数同步失败: {e}"
            return False
        except Exception as e:
            self.get_logger().error(f"参数同步异常 {name}={value}: {e}")
            self.last_error = f"参数同步异常: {e}"
            return False
        finally:
            time.sleep(0.1)
    
    def sync_all_params(self):
        """同步所有参数到摄像头"""
        self.get_logger().info('正在同步所有参数到摄像头...')
        self.sync_param('led_intensity', self.get_parameter('led_intensity').value)
        self.sync_param('quality', self.get_parameter('quality').value)
        self.sync_param('framesize', self.get_parameter('framesize').value)
        self.sync_param('hmirror', self.get_parameter('hmirror').value)
        self.sync_param('vflip', self.get_parameter('vflip').value)
        
        # 人脸检测在某些ESP32-CAM固件上会返回500错误，跳过以避免中断
        face_detect_val = self.get_parameter('face_detect').value
        if face_detect_val == 0:
            self.get_logger().info('跳过face_detect参数同步（已知在某些固件上不稳定）')
        else:
            self.sync_param('face_detect', face_detect_val)
        
        self.sync_param('special_effect', self.get_parameter('special_effect').value)
        self.get_logger().info('参数同步完成')
    
    def wait_connect(self):
        """等待摄像头UDP广播连接"""
        udp_port = self.get_parameter('udp_port').value
        stream_port = self.get_parameter('stream_port').value
        reconnect_timeout = self.get_parameter('reconnect_timeout').value
        
        udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_socket.bind(('', udp_port))
        udp_socket.settimeout(5.0)  # 设置超时避免阻塞
        
        self.get_logger().info(f'UDP监听已启动，端口: {udp_port}')
        
        while rclpy.ok():
            try:
                data, addr = udp_socket.recvfrom(1024)
                
                # 如果超过指定时间没有收到图像，或者是首次连接，重新连接
                time_since_last_frame = time.time() - self.last_recv_img_time
                should_reconnect = (
                    self.last_recv_img_time == 0 or 
                    time_since_last_frame >= reconnect_timeout
                )
                
                if should_reconnect:
                    with self.state_lock:
                        # 检查是否已有活跃的流线程
                        if self.stream_thread and self.stream_thread.is_alive():
                            self.get_logger().info('已有活跃的流连接，跳过重连')
                            continue
                        
                        self.last_recv_img_time = time.time()
                        self.camera_ip = addr[0]
                        self.url = f'http://{self.camera_ip}:{stream_port}/stream'
                        self.camera_state = CameraState.CONNECTING
                        self.reconnect_count += 1
                        
                        self.get_logger().info(f'检测到摄像头: {self.camera_ip}')
                        self.get_logger().info(f'开始读取图像流: {self.url}')
                        self.get_logger().info(f'重连次数: {self.reconnect_count}')
                        
                        # 同步参数
                        self.sync_all_params()
                        
                        # 启动图像读取线程
                        self.stream_thread = threading.Thread(
                            target=self.stream_image, daemon=True
                        )
                        self.stream_thread.start()
                    
            except socket.timeout:
                # 超时是正常的，继续等待
                if self.camera_state == CameraState.DISCOVERING:
                    self.get_logger().debug('等待摄像头UDP广播...')
                continue
            except Exception as e:
                self.get_logger().error(f'UDP接收错误: {e}')
                self.last_error = f"UDP接收错误: {e}"
                self.camera_state = CameraState.ERROR
                time.sleep(1)
    
    def process_images(self):
        """异步处理图像队列（独立线程）"""
        self.get_logger().info('图像处理线程已启动')
        
        while rclpy.ok():
            try:
                # 从队列获取JPEG数据
                jpg_data, timestamp = self.image_queue.get(timeout=1.0)
                
                # 直接发布压缩图像，无需解码
                ros_image = CompressedImage()
                ros_image.header.stamp = timestamp
                ros_image.header.frame_id = "camera_frame"
                ros_image.format = "jpeg"
                ros_image.data = jpg_data
                self.publisher_.publish(ros_image)
                    
            except Exception as e:
                # 队列超时是正常的，继续等待
                if "Empty" not in str(type(e)):
                    self.get_logger().error(f'图像处理异常: {e}')
                continue
        
        self.get_logger().info('图像处理线程已退出')
    
    def stream_image(self):
        """从HTTP流读取图像（优化版）"""
        stream_timeout = self.get_parameter('stream_timeout').value
        enable_frame_skip = self.get_parameter('enable_frame_skip').value
        max_buffer_size = self.get_parameter('max_buffer_size').value
        read_chunk_size = self.get_parameter('read_chunk_size').value
        
        try:
            self.stream = urllib.request.urlopen(self.url, timeout=stream_timeout)
            self.stream_bytes = b''
            self.count = 0
            self.last_time = time.time()
            self.last_recv_img_time = time.time()
            
            with self.state_lock:
                self.camera_state = CameraState.CONNECTED
            
            self.get_logger().info('图像流已连接，开始接收图像...')
            self.get_logger().info(f'帧跳过: {"启用" if enable_frame_skip else "禁用"}')
            self.get_logger().info(f'缓冲区大小: {max_buffer_size} 字节, 读取块: {read_chunk_size} 字节')
            
            while rclpy.ok():
                # 读取数据块 - 使用参数化的块大小
                chunk = self.stream.read(read_chunk_size)
                if not chunk:
                    self.get_logger().warn('图像流中断')
                    with self.state_lock:
                        self.camera_state = CameraState.DISCONNECTED
                    break
                
                self.stream_bytes += chunk
                
                # 防止缓冲区无限增长
                if len(self.stream_bytes) > max_buffer_size:
                    self.get_logger().warn(f'缓冲区溢出 ({len(self.stream_bytes)} 字节)，清空缓冲区')
                    self.stream_bytes = b''
                    self.dropped_frames += 1
                    continue
                
                # 查找JPEG标记
                first = self.stream_bytes.find(b'\xff\xd8')  # JPEG起始标记
                last = self.stream_bytes.find(b'\xff\xd9')   # JPEG结束标记
                
                if first != -1 and last != -1:
                    if last < first:
                        # 结束标记在起始标记之前，丢弃损坏的数据
                        self.stream_bytes = self.stream_bytes[last + 2:]
                        self.dropped_frames += 1
                    else:
                        # 提取完整的JPEG图像
                        jpg = self.stream_bytes[first:last + 2]
                        self.stream_bytes = self.stream_bytes[last + 2:]
                        
                        self.count += 1
                        self.total_frames += 1
                        self.last_recv_img_time = time.time()
                        timestamp = self.get_clock().now().to_msg()
                        
                        # 尝试将JPEG数据放入队列
                        try:
                            if enable_frame_skip:
                                # 非阻塞模式：队列满时跳过帧
                                self.image_queue.put_nowait((jpg, timestamp))
                            else:
                                # 阻塞模式：等待队列有空间
                                self.image_queue.put((jpg, timestamp), timeout=0.1)
                        except Full:
                            # 队列满，跳过此帧
                            self.skipped_frames += 1
                            if self.skipped_frames % 10 == 0:
                                self.get_logger().warn(
                                    f'处理队列已满，已跳过 {self.skipped_frames} 帧'
                                )
                        
                        # 计算帧率
                        if int(time.time() - self.last_time) >= 2:
                            self.rate = self.count / (time.time() - self.last_time)
                            queue_size = self.image_queue.qsize()
                            self.get_logger().info(
                                f"接收帧率: {self.rate:.2f} FPS | "
                                f"总帧数: {self.total_frames} | "
                                f"丢帧: {self.dropped_frames} | "
                                f"跳帧: {self.skipped_frames} | "
                                f"队列: {queue_size}/20"
                            )
                            self.last_time = time.time()
                            self.count = 0
                            
        except urllib.error.URLError as e:
            self.get_logger().error(f'HTTP流连接失败: {e}')
            self.last_error = f"HTTP流连接失败: {e}"
            with self.state_lock:
                self.camera_state = CameraState.ERROR
        except Exception as e:
            self.get_logger().error(f'图像流处理异常: {e}')
            self.last_error = f"图像流处理异常: {e}"
            with self.state_lock:
                self.camera_state = CameraState.ERROR
        finally:
            if self.stream:
                self.stream.close()
            with self.state_lock:
                if self.camera_state == CameraState.CONNECTED:
                    self.camera_state = CameraState.DISCONNECTED
            self.get_logger().info('图像流已关闭')


    def reconnect_callback(self, request, response):
        """手动触发重连服务回调"""
        self.get_logger().info('收到手动重连请求')
        
        try:
            # 重置状态
            with self.state_lock:
                self.last_recv_img_time = 0
                self.camera_state = CameraState.DISCOVERING
            
            response.success = True
            response.message = '已触发重连，等待摄像头UDP广播...'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = f'重连失败: {e}'
            self.get_logger().error(response.message)
        
        return response
    
    def get_status_callback(self, request, response):
        """获取摄像头状态服务回调"""
        with self.state_lock:
            state = self.camera_state
        
        drop_rate = (self.dropped_frames / self.total_frames * 100) if self.total_frames > 0 else 0
        skip_rate = (self.skipped_frames / self.total_frames * 100) if self.total_frames > 0 else 0
            
        status_info = {
            'state': state.name,
            'camera_ip': self.camera_ip or 'N/A',
            'frame_rate': f'{self.rate:.2f} FPS',
            'total_frames': str(self.total_frames),
            'dropped_frames': f'{self.dropped_frames} ({drop_rate:.1f}%)',
            'skipped_frames': f'{self.skipped_frames} ({skip_rate:.1f}%)',
            'queue_size': f'{self.image_queue.qsize()}/5',
            'reconnect_count': str(self.reconnect_count),
            'last_error': self.last_error or 'None'
        }
        
        response.success = True
        response.message = '\n'.join([f'{k}: {v}' for k, v in status_info.items()])
        
        return response
    
    def publish_diagnostics(self):
        """发布诊断信息"""
        if not self.get_parameter('enable_diagnostics').value:
            return
        
        msg = DiagnosticArray()
        msg.header.stamp = self.get_clock().now().to_msg()
        
        status = DiagnosticStatus()
        status.name = f'{self.get_name()}: ESP32 Camera'
        status.hardware_id = self.camera_ip or 'unknown'
        
        # 设置状态级别
        with self.state_lock:
            state = self.camera_state
        
        if state == CameraState.CONNECTED:
            status.level = DiagnosticStatus.OK
            status.message = 'Camera connected and streaming'
        elif state == CameraState.CONNECTING:
            status.level = DiagnosticStatus.WARN
            status.message = 'Connecting to camera...'
        elif state == CameraState.DISCOVERING:
            status.level = DiagnosticStatus.WARN
            status.message = 'Discovering camera via UDP...'
        elif state == CameraState.DISCONNECTED:
            status.level = DiagnosticStatus.ERROR
            status.message = 'Camera disconnected'
        else:  # ERROR
            status.level = DiagnosticStatus.ERROR
            status.message = f'Camera error: {self.last_error or "Unknown"}'
        
        # 添加详细信息
        status.values.append(KeyValue(key='state', value=state.name))
        status.values.append(KeyValue(key='camera_ip', value=self.camera_ip or 'N/A'))
        status.values.append(KeyValue(key='frame_rate', value=f'{self.rate:.2f}'))
        status.values.append(KeyValue(key='total_frames', value=str(self.total_frames)))
        status.values.append(KeyValue(key='dropped_frames', value=str(self.dropped_frames)))
        status.values.append(KeyValue(key='skipped_frames', value=str(self.skipped_frames)))
        status.values.append(KeyValue(key='queue_size', value=str(self.image_queue.qsize())))
        status.values.append(KeyValue(key='reconnect_count', value=str(self.reconnect_count)))
        
        if self.last_error:
            status.values.append(KeyValue(key='last_error', value=self.last_error))
        
        msg.status.append(status)
        self.diagnostics_pub.publish(msg)


def main(args=None):
    rclpy.init(args=args)
    node = ESP32CameraNode()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
